import random
import json
from net.gpt_simple_tcp_server import main as net_run
from net.gpt_simple_tcp_server import Flag
from net.gpt_simple_tcp_server import _Flag
# from net import gpt_simple_tcp_server
from rl_envs.Car2dEnv import Car2dEnv
import threading
from net.msg_handle import msg_handle
from bdtime import tt
import gym
from tools.my_log import log
# print = log.info


class Policy:
    """
    人工控制
    """

    action_conv_dc = {
        'exit': -1,
        'e': -1,
        '0': 0,
        'w': 1,
        's': 2,
        'a': 3,
        'd': 4,
    }

    class ChooseActionPolicy:
        sample = 0
        human = 1
        dqn = 2

    model = ChooseActionPolicy.sample
    # model = ChooseActionPolicy.human

    def _choose_action_by_human(self, state):
        _action = input("输入[wsad], 对应[上下左右]. 输入[exit]退出 >>>")
        action = self.action_conv_dc.get(_action.lower())
        return action

    def _choose_action_by_sample(self, state=None):
        action = random.randint(0, 4)
        return action

    def choose_action(self, state=None):
        action = 0
        if self.model == self.ChooseActionPolicy.sample:
            action = self._choose_action_by_sample(state)
        elif self.model == self.ChooseActionPolicy.human:
            action = self._choose_action_by_human(state)
        return action


policy = Policy()


if __name__ == '__main__':
    # env = Car2dEnv()
    env = gym.make('Car2dEnv-v0')

    # env.action_space
    # env.reset()
    # env.step(1)

    msg_handle.add_listener('MsgTest', env.on_msg_test)
    msg_handle.add_listener('MsgState', env.on_msg_state)
    # net_run()

    # thread_save = threading.Thread(target=gpt_simple_tcp_server.main)
    thread_net = threading.Thread(target=net_run)
    thread_net.setDaemon(True)
    thread_net.start()

    # def choose_action(state):
    #     return

    while msg_handle.client_connection is None:
        if tt.now(0) and tt.now(0) % 5 == 0:
            print('等待客户端连接中...', tt.now(1))
        tt.sleep(1)

    print('--- 已和客户端建立连接!')

    total_timesteps = 100000
    global_step = 0
    current_epoch = 0

    from tqdm import tqdm
    tq = tqdm(total=total_timesteps)
    msg = "初始化tqdm"

    env.reset()
    tt.__init__()
    reward_epoch = 0
    max_reward = -9999999
    max_i = -1
    error_state_counts = 0

    recent_reward_ls = []
    recent_mean_reward = 0      # 最近`recent_mean_len`场平均奖励
    recent_mean_len = 100

    while True:
        if global_step > total_timesteps:
            print('\n*** Main: `global_step`超过最大值`total_timesteps`')
            break

        if not msg_handle.client_connection:
            print('\n*** Main: 客户端连接已断开...')
            break

        if env.i_action < global_step:     # 处理完当前动作再进行下一个
            print('******************', env.i_action, '---', global_step)
            continue

        # _action = input("输入[0-5], 对应[停上下左右] >>>")
        action = policy.choose_action(None)
        # print('=== --- action: ', action)

        if action is not None:
            if action == -1:
                print("--- 输入[exit], 退出游戏, 结束服务端进程.")
                Flag.EXIST_FLAG = True
                break

            observation, r, done, info = env.step(action)
            if observation is not None:
                reward_epoch += r
                if done:
                    reward_epoch = round(reward_epoch, 4)
                    recent_reward_ls.append(reward_epoch)
                    if len(recent_reward_ls) < recent_mean_len:
                        recent_mean_reward = sum(recent_reward_ls) / len(recent_reward_ls)
                    else:
                        recent_mean_reward = sum(recent_reward_ls[-recent_mean_len:]) / recent_mean_len
                    recent_mean_reward = round(recent_mean_reward, 4)
                    if reward_epoch > max_reward:
                        max_reward = reward_epoch
                        max_i = current_epoch

                    msg = f"current_e[{current_epoch}], r: {reward_epoch} |" \
                          f" recent_mean_r[{recent_mean_len}]: {recent_mean_reward} |" \
                          f" max_r: {max_reward}, max_i: {max_i} |" \
                          f" e: {error_state_counts}"

                    reward_epoch = 0
                    current_epoch += 1
            else:
                msg = f"*** ret is None!"
                error_state_counts += 1
            tq.desc = msg

            tq.update(1)
            global_step += 1
        else:
            print(f"--- 输入错误... 输入[wsad]移动, [0]退出...")
    # exit()
    print(f'速度: {round(tq.n / tt.now(), 3)} 条/秒')
    msg_handle.close()

    from bdtime import tt
    import sys

    tt.sleep(1)

    if thread_net.is_alive():
        print('等待线程thread_net释放...')
        tt.sleep(3)
        if thread_net.is_alive():
            print('*** 无法自动释放, 强制结束线程!')
            sys.exit()

